# 构建⼀个带有Human-In-Loop的图
from operator import add
from langchain_core.messages import AnyMessage
from langchain_openai import ChatOpenAI
from langgraph.checkpoint.memory import InMemorySaver
from langgraph.constants import START, END
from langgraph.graph import StateGraph
from typing import Literal, TypedDict, Annotated
from langgraph.types import interrupt, Command
from langchain_core.messages import HumanMessage

# 构建deepseek大模型客户端
# model: 指定使用的大模型为deepseek-chat
# base_url: 指定API的基础URL
# openai_api_key: 指定API密钥
llm = ChatOpenAI(
    model="deepseek-chat",
    base_url="https://api.deepseek.com/v1",
    openai_api_key="sk-ddad02a64c4c47ea9f6f526ef47cb602",
)


class State(TypedDict):
    """
    图中状态的类型定义，用于在节点之间传递状态信息。

    Attributes:
        messages (Annotated[list[AnyMessage], add]): 消息列表，使用add函数进行合并。
    """
    messages: Annotated[list[AnyMessage], add]


def human_approval(state: State) -> Command[Literal["call_llm", END]]:
    """
    人类审批节点函数。该函数会中断执行流程并等待用户输入是否同意调用大语言模型。

    Args:
        state (State): 当前图的状态，包含消息列表等信息。

    Returns:
        Command[Literal["call_llm", END]]: 返回一个命令对象，决定下一步是调用大模型还是结束流程。
    """
    # 中断执行并等待用户输入
    # question: 提示用户的问题
    is_approved = interrupt(
        {
            "question": "是否同意调⽤⼤语⾔模型？"
        }
    )
    if is_approved:
        # 用户同意，继续执行调用大模型节点
        return Command(goto="call_llm")
    else:
        # 用户不同意，结束流程
        return Command(goto=END)


def call_llm(state: State):
    """
    调用大语言模型的节点函数。

    Args:
        state (State): 当前图的状态，包含需要处理的消息列表。

    Returns:
        dict: 包含模型响应结果的消息字典。
    """
    # 调用大语言模型处理消息
    response = llm.invoke(state["messages"])
    return {"messages": [response]}


# 创建状态图构建器，用于定义工作流的状态转换
# State: 状态类，定义了工作流中各个节点间传递的状态数据结构
builder = StateGraph(State)

# 添加工作流节点
# "human_approval": 节点名称，human_approval: 处理_human_approval_节点逻辑的函数
# "call_llm": 节点名称，call_llm: 处理_call_llm_节点逻辑的函数
builder.add_node("human_approval", human_approval)
builder.add_node("call_llm", call_llm)

# 定义工作流的起始边，从START节点连接到human_approval节点
# START: 工作流的起始节点常量
# "human_approval": 目标节点名称
builder.add_edge(START, "human_approval")

# 创建内存存储器，用于保存工作流的检查点状态
# InMemorySaver: 基于内存的检查点存储实现
checkpointer = InMemorySaver()

# 编译状态图构建器生成可执行的工作流图
# checkpointer: 检查点存储器，用于持久化工作流状态
# 返回值: 编译后的工作流图对象
graph = builder.compile(checkpointer=checkpointer)

# 提交任务，等待确认
# configurable: 可配置参数，包含线程ID
# thread_id: 线程ID，用于区分不同的执行线程
thread_config = {"configurable": {"thread_id": 1}}
# 调用图的invoke方法执行任务
# messages: 包含用户消息的列表
# config: 配置参数
res = graph.invoke({"messages": [HumanMessage("湖南的省会是哪⾥？")]}, config=thread_config)
# 执⾏后会中断任务，等待确认
print("res:", res)

# 确认同意，继续执⾏任务
# final_result = graph.invoke(Command(resume=True), config=thread_config)
# print("final_result:", final_result)
# 不同意，终⽌任务
# Command(resume=False): 指示不继续执行任务
final_result = graph.invoke(Command(resume=False), config=thread_config)
print("final_result:", final_result)
